### TODO

# Convert a vector with values for each voxel into a 3D array to be saved as a nifti file
## Arguments
## - vec.data: vector with data
## - coord.mat: Matrix of coordinates to transplant into 3D
## - ref.img: reference image (uses it to get array dimensions)
## - use.rownames: TRUE or FALSE (default=TRUE), rownames have numbers to indicate vector number of voxel in 3D => use them
## Return:
## - 3D array
ni.1Dto3D = function(vec.data, coord.mat, ref.img, use.rownames=TRUE) {
	# Number of voxels with values
	num.vox = nrow(coord.mat)
	
	# Dimensions
	ref.dim = dim(ref.img)
	
	# Create the new matrix
	new.mat = array(0, ref.dim)
        
	# Use rownames to identify voxel coordinate?
	if (use.rownames)
		coord.vec = as.numeric(rownames(coord.mat))
        
	# Loop through and save data
	for (i in 1:num.vox) {
		if (use.rownames) {
			## Quick copy
			new.mat[coord.vec[i]] = vec.data[i]
		} else {
			## Get coordinates
			x = coord.mat[i,1]
			y = coord.mat[i,2]
			z = coord.mat[i,3]
			## Save
			new.mat[x,y,z] = vec.data[i]
		}
	}
	
	return(new.mat)
}

ni.write.vector = function(vec, ref.file, output.file, check.output=1) {
    num.regions = length(vec)
    
    # Read in reference image
    ref.img = nifti.image.read(ref.file)
    ref.coords = ni.get.coords(ref.img)
    ref.img[,,] = 0
    
    # Convert vec to 3D matrix
    tmp.3D = ni.1Dto3D(vec, ref.coords, ref.img)
    
    # Save into nifti image
    ref.img[,,] = tmp.3D
    nifti.set.filenames(ref.img, output.file, check=check.output)
    nifti.image.write(ref.img)
}

# Read in functional image into neurots format
# (basically convert a 4D array to a 2D one)
#
# Input Arguments
#   - func: The main 4D functional image.
#       Can Be:
#       1. nifti filename
#       2. 4D nifti object
#       3. 4D array
#   - mask: Voxels to use/include from 4D functional image
#       Can Be:
#       1. nifti filename
#       2. vector of voxels to include
#       3. coords object where 4th 'val' column specifies voxels to include
#       4. 3D array or nifti object (which is used to create a coords object)
#   - ref: 3D nifti reference image (to be used when want to create new output)
#       Can Be:
#       1. nifti filename
#       2. nifti object
#
# Return:
#   - neurots object
neurots <- function(func, mask=NULL, ref=NULL) {
    stime <- proc.time()[3]
    
    # 1. Setup functional image
    if (is.character(func))
        func <- nifti.image.read(abspath(func))
    
    if (class(func) == "nifti")
        fname <- func$fname
    else
        fname <- ""
    
    # 2. Convert 4D to 2D
    if (is.array(func) || class(func) == "nifti")
        mat <- ni4Dto2D(func)
    else
        stop("Unrecognized type/class given for 'func' argument: ", class(func))
    
    # 3. Prepare mask of voxels to use
    if (is.null(mask))
        ## only include voxels with non-zero values
        mask <- apply((func[,,,]>0)*1, c(1,2,3), min)
    else if (is.character(mask))
        mask <- nifti.image.read(abspath(mask))
    if (class(mask) == "nifti" && is.null(ref) == TRUE)
        ref <- mask
    
    # 4. Mask mat and get corresponding 3D coordinates for each voxel if available
    if (is.vector(mask)) {
        coords <- data.frame()
        mat <- mat[,mask]
    }
    else if (is.data.frame(mask)) {
        if (nrow(mask) != ncol(mat))
            stop("If mask is a data frame, it must have same number of voxels or rows as functional image")
        coords <- mask
        mat <- mat[,coords$val]
    }
    else if (is.array(mask) || class(mask) == "nifti") {
        if (length(dim(mask)) != 3)
            stop("If mask is array or nifti object, it must be three dimensional")
        coords <- niGetCoords(mask[,,])
        mat <- mat[,coords$val]
    } else {
        stop("Unrecognized type/class given for 'mask' argument: ", class(mask))
    }
    
    # 5. Deal with reference image (require it to be 3D)
    if (is.null(ref)) {
        ## use 3D aspect of functional image
        ref <- nifti.image.copy.info(func)
        ref$dim <- dim(ref)[1:3]
        nifti.set.filenames(ref, sub(basename(ref$fname), "neurots_ref", ref$fname))
    } else if (is.character(ref)) {
        ref <- nifti.image.read(abspath(ref))
    }
    
    if (length(dim(ref)) != 3)
        stop("'ref' input argument must be three dimensional")
    
    # 6. Create new neurots object
    x <- new("neurots", mat=mat, coords=coords, fname=fname, ref=ref)
    
    return(x)
}


# Convert 4D fMRI dataset to a 2D matrix
#
# Arguments
#   - img: Your 4D functional image (where the 4th dimension is time)
#
# Results
#   - matrix of time-points (rows) x voxels (columns)
ni4Dto2D = function(img) {
    # Read in file
    if (is.character(img))
        img <- nifti.image.read(abspath(img))

    # Get image dimensions (assume that 4D)
    imgDim = dim(img)
    if (length(imgDim) != 4)
        stop("ni4Dto2D needs input to be four dimensional")

    # If nifti class (from Rniftilib) then convert to array
    # this is to allow reshape step to work
    if (class(img) == "nifti")
        img = img[,,,]

    # Number of voxels
    nVoxs = prod(imgDim[1:3])
    
    # Number of time points
    nTimePts = imgDim[4]
    
    # Reshape 4D functional image to 2D matrix
    dim(img) = c(nVoxs, nTimePts)
    img <- aperm(img)
    
    return(img)
}


# Create matrix with list of regions and corresponding location in matrix/array
# If you convert a 3D array into a 1D vector, this function allows you to go back to 3D space
# Also this function is useful to go from 2D matrix to 4D array...etc
#
# Arguments:
#   - img: 3D nifti array
#   - thresh: threshold for which voxels to set as TRUE (see 4th column 'val' descrip below)
#       default = 0 (i.e. set as TRUE all positive voxels)
#       if thresh is NULL, then all voxels have TRUE in 4th 'val' column
#
# Return:
#   - matrix of coordinates
#   - each row is a voxel
#   - columns are x, y, and z coordinates
#     4th column is called 'val' and is TRUE or FALSE depending on whether to use this voxel
niGetCoords = function(img, thresh=0) {
    # Read in file
    if (is.character(img))
        img <- nifti.image.read(abspath(img))
    
    # Convert nifti to array
    if (class(img) == "nifti")
        img <- img[,,]
    
    # Get list of indices in each dimension
    indices <- lapply(dim(img), seq)
    
    # Get coordinates for each voxel
    coords <- expand.grid(indices)
    
    # Get only non-zero coordinates in img
    if (is.null(thresh))
        coords <- cbind(coords, rep(TRUE, nrow(coords)))
    else
        coords <- cbind(coords, as.vector(img>thresh))
    
    colnames(coords) <- c("x", "y", "z", "val")
    
    return(coords)
}

###
# GETTERS
###

setMethod("[",
    signature(x = "neurots", drop="missing"),
    function(x, i, j) {
        return(x@mat[i,j])
    }
)

setMethod("[",
    signature(x = "neurots", i="missing", drop="missing"),
    function(x, j, drop=False) {
        return(x@mat[,j])
    }
)

setMethod("[",
    signature(x = "neurots", j="missing", drop="missing"),
    function(x, i) {
        return(x@mat[i,])
    }
)

setMethod("[",
    signature(x = "neurots", i="missing", j="missing", drop="missing"),
    function(x) {
        return(x@mat)
    }
)


###
# SETTERS
###

setMethod("[<-",
    signature(x = "neurots"),
    function(x, i, j, value) {
        x@mat[i,j] <- value
    }
)

setMethod("[<-",
    signature(x = "neurots", i="missing"),
    function(x, j, value) {
        x@mat[,j] <- value
    }
)

setMethod("[<-",
    signature(x = "neurots", j="missing"),
    function(x, i, value) {
        x@mat[i,] <- value
    }
)

setMethod("[<-",
    signature(x = "neurots", i="missing", j="missing"),
    function(x, value) {
        x@mat <- value
    }
)

###
# Print/Show Related
###

setMethod('head', signature(x="neurots"),
    function(x, ...) {
        head(x@mat, ...)
    }
)

setMethod('tail', signature(x="neurots"),
    function(x, ...) {
        tail(x@mat, ...)
    }
)

###
# Other
###
setMethod('dim', signature(x='neurots'),
    function(x) dim(x@mat)
)

setMethod('ncol', signature(x='neurots'),
    function(x) ncol(x@mat)
)

setMethod('nrow', signature(x='neurots'),
    function(x) nrow(x@mat)
)